#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : mngs_batch_check_gender.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/10/22 17:33
@Description :
批量绘制人类基因组覆盖度（10K窗口）并判断性别。
流程步骤：
1. BWA 比对参考基因组
2. 使用 mosdepth 统计覆盖度
3. 使用 plot_coverage.py 绘制覆盖度分布图
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
from multiprocessing import Pool
from datetime import datetime
import subprocess
import logging


# -------------------------------------------------------------------------------------
# 全局配置
# -------------------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s] %(levelname)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
LOGGER = logging.getLogger("CHECK_GENDER")

# 工具路径
PYTHON3 = Path("/home/bioinfo/software/miniconda3/envs/python-venv/bin/python3")
PLOT_PY = Path("/home/bioinfo/repositories/bjzd-RNA-seq/scripts/plot_coverage.py")
REF_FNA = Path("/data/jinlj/project/gender/plot_coverage/database/T2T.clean.fna")
BWA = Path("/home/bioinfo/software/miniconda3/envs/bwa/bin/bwa")
SAMTOOLS = Path("/home/bioinfo/software/miniconda3/envs/samtools/bin/samtools")
MOSDEPTH = Path("/home/bioinfo/software/miniconda3/envs/mosdepth/bin/mosdepth")


# -------------------------------------------------------------------------------------
# 辅助函数
# -------------------------------------------------------------------------------------
def _run_or_dead(sh_cmd: str, sh_file: Path, log_file: Path, ok_file: Path):
    """
    执行 shell 命令。若失败则记录错误，成功则创建 .ok 文件。
    """
    if ok_file.exists():
        LOGGER.info(f"跳过 {sh_file}（已完成）")
        return

    try:
        LOGGER.info(f"执行任务：{sh_file}")
        with open(sh_file, "w", encoding="utf-8") as fs:
            fs.write(sh_cmd)

        with open(log_file, "w", encoding="utf-8") as fl:
            result = subprocess.run(
                f"bash {sh_file}",
                shell=True,
                text=True,
                stdout=fl,
                stderr=fl,
            )

        if result.returncode != 0:
            LOGGER.error(f"执行失败：{sh_file}，退出码 {result.returncode}")
            return

        ok_file.touch()
        LOGGER.info(f"完成：{sh_file}")

    except Exception as e:
        LOGGER.error(f"执行出错：{sh_file}，错误：{e}")


def _run_cmds(sh_files: list[Path], sh_cmds: list[str], cpu_n: int = 1):
    """
    并行执行多个命令脚本。
    """
    ok_files = [f.with_suffix(".ok") for f in sh_files]
    log_files = [f.with_suffix(".log") for f in sh_files]

    with Pool(cpu_n) as pool:
        for sh_cmd, sh_file, ok_file, log_file in zip(sh_cmds, sh_files, ok_files, log_files):
            pool.apply_async(_run_or_dead, args=(sh_cmd, sh_file, log_file, ok_file))
        pool.close()
        pool.join()

    for ok_file in ok_files:
        if not ok_file.exists():
            raise RuntimeError(f"缺少完成标志文件：{ok_file}")


# -------------------------------------------------------------------------------------
# 主流程类定义
# -------------------------------------------------------------------------------------
class Base:
    """任务基础类：加载样本信息并准备目录结构"""

    def __init__(self, args: Namespace):
        self.step_n = 0
        self.begin_time = datetime.now()
        self.infile = Path(args.infile).resolve()
        self.outdir = Path(args.outdir).resolve()
        self.logdir = Path(args.logdir).resolve()
        self.cpu_n = int(args.cpu_n)

        # 样本与原始 reads 路径
        self.samples: list[str] = []
        self.raw_reads: list[str] = []

        with open(self.infile, "r", encoding="utf-8") as fr:
            for line in fr:
                if line := line.strip():
                    sample, raw_reads = line.split()
                    self.samples.append(sample)
                    self.raw_reads.append(raw_reads)

    def run(self):
        """创建输出与日志根目录"""
        self.outdir.mkdir(parents=True, exist_ok=True)
        self.logdir.mkdir(parents=True, exist_ok=True)


class MapGenome(Base):
    """步骤1：BWA 比对到人类参考基因组"""

    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.map_step_name = f"step{self.step_n}_map_genome"
        self.map_outdir = self.outdir / self.map_step_name
        self.map_logdir = self.logdir / self.map_step_name
        self.map_out_bams = [self.map_outdir / f"{s}.bam" for s in self.samples]

    def run(self):
        super().run()
        self.map_outdir.mkdir(parents=True, exist_ok=True)
        self.map_logdir.mkdir(parents=True, exist_ok=True)

        _run_cmds(
            [self.map_logdir / f"{s}.sh" for s in self.samples],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -euo pipefail\n\n"
                    f"{BWA} mem -t 8 {REF_FNA} {r} | "
                    f"{SAMTOOLS} view -@ 8 -F 4 -b -h | "
                    f"{SAMTOOLS} sort -@ 8 -o {bam}\n"
                )
                for r, bam in zip(self.raw_reads, self.map_out_bams)
            ],
            self.cpu_n,
        )


class MosDepth(MapGenome):
    """步骤2：使用 mosdepth 计算覆盖度（10K窗口）"""

    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.mos_step_name = f"step{self.step_n}_mos_depth"
        self.mos_outdir = self.outdir / self.mos_step_name
        self.mos_logdir = self.logdir / self.mos_step_name
        self.mos_out_prefix = [self.mos_outdir / s for s in self.samples]
        self.mos_out_beds = [self.mos_outdir / f"{s}.regions.bed.gz" for s in self.samples]

    def run(self):
        super().run()
        self.mos_outdir.mkdir(parents=True, exist_ok=True)
        self.mos_logdir.mkdir(parents=True, exist_ok=True)

        _run_cmds(
            [self.mos_logdir / f"{s}.sh" for s in self.samples],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"set -euo pipefail\n\n"
                    f"{SAMTOOLS} index {bam}\n\n"
                    f"{MOSDEPTH} -n -t 8 --by 10000 {pf} {bam}\n"
                )
                for pf, bam in zip(self.mos_out_prefix, self.map_out_bams)
            ],
            self.cpu_n,
        )


class PlotCoverage(MosDepth):
    """步骤3：绘制覆盖度图像"""

    def __init__(self, args: Namespace):
        super().__init__(args)
        self.step_n += 1
        self.plot_step_name = f"step{self.step_n}_plot_coverage"
        self.plot_outdir = self.outdir / self.plot_step_name
        self.plot_logdir = self.logdir / self.plot_step_name
        self.plot_out_pngs = [self.plot_outdir / f"{s}.png" for s in self.samples]
        self.plot_out_pdfs = [self.plot_outdir / f"{s}.pdf" for s in self.samples]

    def run(self):
        super().run()
        self.plot_outdir.mkdir(parents=True, exist_ok=True)
        self.plot_logdir.mkdir(parents=True, exist_ok=True)

        _run_cmds(
            [self.plot_logdir / f"{s}.sh" for s in self.samples],
            [
                (
                    f"#!/bin/bash\n\n"
                    f"{PYTHON3} {PLOT_PY} "
                    f"--in_bed {bed} "
                    f"--out_png {png} "
                    f"--out_pdf {pdf}\n"
                )
                for bed, png, pdf in zip(self.mos_out_beds, self.plot_out_pngs, self.plot_out_pdfs)
            ],
            self.cpu_n,
        )


class End(PlotCoverage):
    """最终步骤：输出耗时"""

    def __init__(self, args: Namespace):
        super().__init__(args)

    def run(self):
        super().run()
        elapsed = datetime.now() - self.begin_time
        LOGGER.info(f"任务完成，总耗时：{elapsed}")


# -------------------------------------------------------------------------------------
# 主入口
# -------------------------------------------------------------------------------------
if __name__ == "__main__":
    parser = ArgumentParser(description="批量性别判断：比对 → 深度统计 → 覆盖绘图")
    parser.add_argument("--infile", type=str, required=True, help="输入样本信息文件，每行：sample fastq")
    parser.add_argument("--outdir", type=str, required=True, help="输出目录")
    parser.add_argument("--logdir", type=str, required=True, help="日志目录")
    parser.add_argument("--cpu_n", type=int, required=True, help="并行任务数")
    End(parser.parse_args()).run()
